#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    MOTEvaluationNode
【子结点名称】
    MOTLoaderNode
    MOTResultAssembleNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    spirecv_msgs::2DTargets （目标检测结果）
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    mot_folder_path (str): MOT测试文件夹地址
    pred_file_path (str): 评价结果文件夹地址
    mot_path (str): MOT根文件夹地址
    algorithm (str): 使用的目标检测节点
    mot (str): mot版本
【备注】：
    无
"""

import threading
import os
import argparse
from typing import Union
from spirecv.base.BaseNode import BaseNode

from spirecv.dataloader.MOTLoaderNode import MOTLoaderNode
from spirecv.evaluator.MOTResultAssembleNode import MOTResultAssembleNode


class MOTEvaluationNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
    ):
        self.job_name = job_name
        self.param_dict_or_file = param_dict_or_file
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown
        )
        self.mot_path = self.get_param("mot_path", os.path.expanduser('~')+"/mot20/")
        self.mot_folder_path = self.get_param("mot_folder_path", os.path.expanduser('~')+"/mot20/train/")
        self.pred_file_path = self.get_param("pred_file_path", os.path.expanduser('~')+"/mot20/pred_mot20/data/")
        self.algorithm = self.get_param("algorithm", "yolov5v7")
        self.mot = self.get_param("mot", "mot20")
        self.visualizer = self.get_param("visualizer","False")
        self.start()

    def release(self):
        BaseNode.release(self)

    def run(self):
        dataloader = MOTLoaderNode(self.job_name,mot_folder_path=self.mot_folder_path)
        if self.algorithm =='yolov5v7':
            from spirecv.algorithm.common_det.YOLOv5v7DetNode import YOLOv5v7DetNode
            YOLOv5v7DetNode(self.job_name, param_dict_or_file=self.param_dict_or_file)
        MOTResultAssembleNode(self.job_name, mot_folder_path=self.mot_folder_path,\
                              pred_file_path=self.pred_file_path,mot_path=self.mot_path,mot=self.mot)
        dataloader.launch_next()
        if self.visualizer == 'True':
            from spirecv.visualizer.DetectionVisNode import DetectionVisNode
            DetectionVisNode(self.job_name,'/'+self.job_name+'/mot/image_results',imshow=True)
        self.release()
        print('{} quit!'.format(self.__class__.__name__))
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'param_dict_or_file',
        help='param_dict_or_file')
    parser.add_argument(
        '--job_name',
        type=str,
        default='detection_evaluation',
        help='job-name')
    args = parser.parse_args()
    node = MOTEvaluationNode(args.job_name,param_dict_or_file = args.param_dict_or_file)

